#include "Clean.h"
#include "core/float_compare.h"

#include <algorithm>
#include <limits>

#include <iostream>
#include <cstdio>
#include <cmath>

using namespace std;

namespace spectClean {


Clean::Clean()
{
}

Clean::~Clean()
{
}


void Clean::simpleClean(const Spectrum & psf, Spectrum & cleanComponents,
        Spectrum & residual, const unsigned num_iterations, const double gain)
{
    cout << "= psf nfreqs = " << psf.size() << endl;
    cout << "= resid nfreqs = " << residual.size() << endl;

    Complex * res = &residual.amp()[0];
    Complex * cc = &cleanComponents.amp()[0];
    const Complex * w = &psf.amp()[0];
    const unsigned nFreqs = residual.size();
    const unsigned iFreq0 = nFreqs / 2;
    Spectrum temp;
    temp.freq() = residual.freq();
    temp.resize(residual.size());
    Complex * t = &temp.amp()[0];
    int num_rejected_components = 0;

    // TODO: add some kind of rms end condition?
    // i.e. 1. end if RMS gets bigger
    //      2. end if peak = RMS?

    printf("iFreq0 = %i\n", iFreq0);
    for (unsigned j = 0; j < num_iterations; ++j)
    {
        // Find max residual (only searching +ve part of the spectrum).
        //unsigned iMax = findMaxIndex(&res[iFreq0], iFreq0) + iFreq0;
        unsigned iMax = findMaxIndex2(res, iFreq0, nFreqs-1);
        const double wMag = abs(w[2 * iMax]);

        // FIXME: IF THE max is at 0 frequency (i.e. window is pure real)
        // denominator = 0! => nan.
        if (approx_equal(wMag, 1.0))
        {
            cout << "****************************************" << endl;
            cout << "**** eeek, divide by zero coming up! ***" << endl;
            cout << "**** 0 Hz (DC) spectrum peak         ***" << endl;
            cout << "****************************************" << endl;
            printf("bad index = %i %f (%i, %i)\n", iMax, residual.freq()[iMax], iMax - iFreq0, iFreq0);

            ++num_rejected_components;
            break;
        }

        // Find the amplitude taking into account of the contamination of the
        // negative frequency peak.
        const double rRe = res[iMax].real();
        const double rIm = res[iMax].imag();
        const double wRe = w[2 * iMax].real();
        const double wIm = w[2 * iMax].imag();

        const double denominator = 1.0 - pow(wMag, 2);
        const double aRe = (rRe * (1.0 - wRe) - (rIm * wIm)) / denominator;
        const double aIm = (rIm * (1.0 + wRe) - (rRe * wIm)) / denominator;

//        cout << endl;
//        cout << "= clean iteration [" << j << "] " << endl;
//        cout << "     - max = " << abs(res[iMax]) << " at " << iMax
//             << " (" << residual.freq()[iMax] << "  Hz)" <<endl;
//        cout << "     - " << abs(w[2 * iMax]) << endl;
//        cout << "     - psf freq   = " << psf.freq()[2 * iMax] << endl;
//        cout << "     - r freq     = " << residual.freq()[iMax] << endl;
//        cout << "     - r (re, im) = " << rRe << " " << rIm << endl;
//        cout << "     - w (re, im) = " << wRe << " " << wIm << endl;
//        cout << "     - a (re, im) = " << aRe << " " << aIm << endl;


        // Append the clean component.
        cc[iMax] += Complex(gain * aRe, gain * aIm);
//        residual.save("test_clean_residual_" + QString::number(j) + ".txt");

        // Update the residual spectrum by subtracting the psf.
        double tempMax = numeric_limits<double>::min();
        for (unsigned i = 0; i < nFreqs; ++i)
        {
            const Complex w1 = w[i - iMax + (nFreqs - 1)];
            const Complex w2 = w[i + iMax];
            double re = aRe * (w1.real() + w2.real()) + aIm * (w2.imag() - w1.imag());
            double im = aRe * (w1.imag() + w2.imag()) + aIm * (w1.real() - w2.real());
            re *= gain;
            im *= gain;
            t[i] = Complex(re, im);
            res[i] -= Complex(re, im);
            tempMax = max(re * re + im * im, tempMax);
        }
//        temp.save("test_temp_" + QString::number(j) + ".txt");
//        cout << "tempMax = " << tempMax << endl;

        // TODO: get the residual rms?!
        // TODO: save the rms and peak value / position per iteration.
        // as this will allow to follow the clean.
    }
}


unsigned Clean::findMaxIndex(const Complex * data, const unsigned n)
{
    double max = -numeric_limits<double>::max();
    unsigned maxIdx = 0;
    for (unsigned i = 0; i < n; ++i)
    {
        const double re = data[i].real();
        const double im = data[i].imag();
        const double amp = re * re + im * im;
        if (amp > max)
        {
            max = amp;
            maxIdx = i;
        }
    }
    return maxIdx;
}


unsigned Clean::findMaxIndex2(const Complex * data, const unsigned index_start,
        const unsigned index_end, const double old_max)
{
    unsigned maxIdx = 0;
    double max = old_max;
    for (unsigned i = index_start; i <= index_end; ++i)
    {
        const double re = data[i].real();
        const double im = data[i].imag();
        const double amp = re * re + im * im;
        if (amp > max)
        {
            max = amp;
            maxIdx = i;
        }
    }
    return maxIdx;
}



void Clean::make_clean_spectrum(const Spectrum & clean_components,
        Spectrum & clean_spectrum)
{
    const int num_freqs = clean_components.size();
    clean_spectrum.resize(num_freqs);
    const Complex * cc = &clean_components.amp()[0];
    Complex * cs = &clean_spectrum.amp()[0];
    clean_spectrum.freq() = clean_components.freq();

    // TODO: need to set this properly from the window function width.
    int radius = 10;
    vector<double> cb(2 * radius + 1, 0.0);
    for (int i = -radius; i <= radius; ++i)
    {
        double arg = -double(i * i) / 2.0 * double(0.4);
        cb[i + radius] = exp(arg);
    }


    for (int i = 0; i < num_freqs; ++i)
    {
        double re = 0.0;
        double im = 0.0;

//        printf("%i\n", i);
        for (int k = -radius; k <= radius; ++k)
        {
            int index = i + k;
            if (index < 0 || index >= num_freqs)
                continue;
//            printf("   %i\n", index);
            re += cc[index].real() * cb[k + radius];
            im += cc[index].imag() * cb[k + radius];
        }
        cs[i] = Complex(re, im);
    }
}






} // namespace spectClean










///**
// * @details Makes the Gaussian clean beam
// */
//void Clean::makeCleanBeam(SpectrumData * window, SpectrumData * cleanBeam)
//{
//    timeval start, end;
//    gettimeofday(&start, 0);
//
//    // Get the FWHM of the window function peak.
//    double fwhm = 0.0;
//    unsigned iFwhm = 0;
//    _returnFwhmWindow(window, &fwhm, &iFwhm);
//
//    // Set various clean beam spectrum parameters.
//    cleanBeam->_freqInc = window->_freqInc;
//    cleanBeam->_type = SPECT_CLEAN_BEAM;
//
//    // Number of frequencies over which to generate the clean beam
//    unsigned support = 5;
//    int freqZeroWindow = window->_freq.size() / 2;
//    int width = support * (iFwhm - freqZeroWindow);
//    int nFreqs = 2 * width + 1;
//    int iFreqZero = nFreqs / 2;
//
//    cleanBeam->_rangeFreq._min = 0.0;
//    cleanBeam->_rangeFreq._max = width * cleanBeam->_freqInc;
//
//    // Resize the spectrum data arrays.
//    cleanBeam->_freq.resize(nFreqs);
//    cleanBeam->_cAmp.resize(nFreqs);
//
//    // Calculate the variance of the Gaussian clean beam.
//    double variance = std::pow(fwhm / (2.0 * std::sqrt(2.0 * std::log(2.0))), 2);
//
//    QString title = "Making clean beam";
//    title += " (" + QString::number(nFreqs) + " points)";
//
//    // Generate the clean beam values.
//    #pragma omp parallel for
//    for (int f = 0; f < nFreqs; ++f)
//    {
//        double freq = double(f - iFreqZero) * cleanBeam->_freqInc;
//        double amp = exp(- (pow(freq, 2) / (2.0 * variance)));
//        cleanBeam->_freq[f] = freq;
//        cleanBeam->_cAmp[f] = amp;
//    }
//
//    gettimeofday(&end, NULL);
//    double timeTaken = (end.tv_sec - start.tv_sec) +
//            (end.tv_usec - start.tv_usec) / 1.0e6;
//
//    printf("- %-30s [Time: %.2f s]\n","Make clean beam", timeTaken);
//    printf("- stdev: %f\n", sqrt(variance));
//}
//
//
///**
// * @details
// */
//void Clean::initClean(SpectrumData * dirty, SpectrumData * residual,
//        SpectrumData * components, SpectrumData * clean,
//        SpectrumData * cleanResidual)
//{
//    if (!clean->_amp.empty())
//        return;
//
//    cout << "= Initialising spectra for clean" << endl;
//
//    unsigned nFreqs = dirty->_freq.size();
//    double freqInc = dirty->_freqInc;
//    double minFreq = dirty->_rangeFreq._min;
//    double maxFreq = dirty->_rangeFreq._max;
//
//    Range range(minFreq, maxFreq);
//    residual->init(SPECT_RESIDUAL, nFreqs, range, freqInc);
//    components->init(SPECT_COMPONENTS, nFreqs, range, freqInc);
//    clean->init(SPECT_RESIDUAL, nFreqs, range, freqInc);
//    cleanResidual->init(SPECT_RESIDUAL, nFreqs, range, freqInc);
//
//    residual->_cAmp = dirty->_cAmp;
//
//    clean->_freq = dirty->_freq;
//    cleanResidual->_freq = dirty->_freq;
//    components->_freq = dirty->_freq;
//    residual->_freq = dirty->_freq;
//}
//
//
///**
// * @detail Clean the dirty spectrum
// */
//void Clean::cleanSpectrum(unsigned loops, double gain, SpectrumData * window,
//        SpectrumData * residual, SpectrumData * components)
//{
//    timeval start;
//    timeval end;
//    gettimeofday(&start, NULL);
//
//    // The number of frequencies in the residual spectrum.
//    const int nFreqs = residual->_freq.size();
//
//    // Index of the zero frequency.
//    // note: this assumes the number of frequencies is even?
//    const int iFreq0 = nFreqs / 2;
//
//    // Counter for the number of clean components found
//    int iCleanComponent = 0;
//
//    // Counter for the number of bad indices found
//    int iBad = 0;
//
//    // Clean loop
//    for (unsigned iLoop = 0; iLoop < loops; ++iLoop)
//    {
//        // Find the max amp in the residual spectrum
//        unsigned iMax = _findMaxIndex(residual->_cAmp, iFreq0, (nFreqs-1));
//
//        //  If the index is bad (divide by zero problem)
//        if (approxEqual(abs(window->_cAmp[2 * iMax]), 1.0))
//        {
//            // Increment the bad index counter and print a warning message.
//            iBad++;
//
//            printf("Peak at residual[%d]: freq = %f, amp = %f\n",
//                    iMax, residual->_freq[iMax], abs(residual->_cAmp[iMax]));
//            printf("WARNING: [loop %d] Bad index. Finding a new max value.\n",
//                    iLoop);
//
//            // Find a few max residual
//            iMax = _findMaxSafeIndex(residual->_cAmp, iFreq0, (nFreqs-1), iMax);
//
//            printf("New peak at residual[%d]: freq = %f, amp = %f\n",
//                    iMax, residual->_freq[iMax], abs(residual->_cAmp[iMax]));
//        }
//
//        // Set the values of the peak in the residual spectra.
//        double maxFreq = residual->_freq[iMax];
//        double maxAmp = abs(residual->_cAmp[iMax]);
//
//        // Check for a divide by zero error again and break out of the clean
//        // if required
//        if (approxEqual(abs(window->_cAmp[2 * iMax]), 1.0))
//        {
//            printf("ERROR: value in window spectrum would cause a divide by "
//                    "zero error\n");
//            break;
//        }
//
//        // Compute the amplitude of the peak component.
//        double rRe = residual->_cAmp[iMax].real();
//        double rIm = residual->_cAmp[iMax].imag();
//        double wRe = window->_cAmp[2 * iMax].real();
//        double wIm = window->_cAmp[2 * iMax].imag();
//        double wMag = abs(window->_cAmp[2 * iMax]);
//
//        // Get the spectral amp. corresponding to iMax taking into account
//        // of the contamination of the -ve peak.
//        double aRe = ( (rRe * (1.0 - wRe)) - (rIm * wIm) ) / (1.0 - pow(wMag, 2));
//        double aIm = ( (rIm * (1.0 + wRe)) - (rRe * wIm) ) / (1.0 - pow(wMag, 2));
//
//        // Generate the new residual spectrum.
//        for (int f = 0; f < nFreqs; ++f)
//        {
//            double subRe = 0.0;
//            double subIm = 0.0;
//            Complex w1 = window->_cAmp[f - iMax + (nFreqs-1)];
//            Complex w2 = window->_cAmp[f + iMax];
//
//            subRe = aRe * (w1.real() + w2.real()) + aIm * (w2.imag() - w1.imag());
//            subIm = aRe * (w1.imag() + w2.imag()) + aIm * (w1.real() - w2.real());
//            subRe *= gain;
//            subIm *= gain;
//            residual->_cAmp[f] -= Complex(subRe, subIm);
//        }
//
//        // Add the clean component.
//        components->_cAmp[iMax] += Complex(gain * aRe, gain * aIm);
//
//        // Save the clean history.
//        _cleanHistory.peakAmp.push_back(maxAmp);
//        _cleanHistory.peakFreq.push_back(maxFreq);
//        _cleanHistory.extraComponentAmp.push_back(abs(Complex(gain*aRe, gain*aIm)));
//        _cleanHistory.totalComponentAmp.push_back(abs(components->_cAmp[iMax]));
//
//        // Increment the clean component counter.
//        iCleanComponent++;
//    }
//
//    gettimeofday(&end, 0);
//    double timeTaken = (end.tv_sec - start.tv_sec) +
//            (end.tv_usec - start.tv_usec) / 1.0e6;;
//    printf("= %-30s [Time: %.2f s]\n","Clean complete", timeTaken);
//    printf("---------------------------------------------------------------\n");
//    printf("\n");
//}
//
//
//
//
///**
// * @details
// * Construct the clean spectrum from the clean beam and clean components
// *
// * @return an error code
// */
//void Clean::makeCleanSpect(SpectrumData * clean, SpectrumData * cleanResidual,
//        SpectrumData *components, SpectrumData * residual, SpectrumData * cleanBeam)
//{
//    // Get the number of frequencies.
//    int nFreqs = components->_freq.size();
//    int cFreqs = cleanBeam->_freq.size();
//    int cZero = cFreqs / 2;
//
//    for (int f = 0; f < nFreqs; ++f)
//    {
//        double sumRe = 0.0;
//        double sumIm = 0.0;
//
//        for (int b = 0; b < cFreqs; b++)
//        {
//            if ((f-cZero+b >= 0) && (f-cZero+b < nFreqs))
//            {
//                double cbAmp = abs(cleanBeam->_cAmp[b]);
//                int i = f - cZero + b;
//                sumRe += cbAmp * components->_cAmp[i].real();
//                sumIm += cbAmp * components->_cAmp[i].imag();
//            }
//        }
//
//        double rRe = residual->_cAmp[f].real();
//        double rIm = residual->_cAmp[f].imag();
//
//        clean->_cAmp[f] = Complex(sumRe, sumIm);
//        cleanResidual->_cAmp[f] = Complex(sumRe + rRe, sumIm + rIm);
//    }
//}
//
//
//
//
//
//////////////////////////////////////////////////////////////////////////////////
//// Private members
//////////////////////////////////////////////////////////////////////////////////
//
///**
// * @details
// * Find the full width half maximum of central peak
// * in the window function.
// *
// * @param[in]  window 	Window spectrum
// * @param[out] fwhm 	FWHM of the window spectrum peak
// * @param[out] iFwhm	Index of the frequency at the FWHM
// *
// * @return the fwhm of the window spectrum
// */
//void Clean::_returnFwhmWindow(SpectrumData * window, double * fwhm,
//        unsigned * iFwhm)
//{
//    // Store the maximum in the window function (assumed to be at 0 freq)
//    int iZeroFreq = window->_freq.size() / 2;
//
//    double max = std::abs(window->_cAmp[iZeroFreq]);
//    double halfMax = max / 2.0;
//
//    // Find the frequency index where the amp drops to half the maximum
//    unsigned int iFreq = iZeroFreq;
//    while (std::abs(window->_cAmp[iFreq]) > halfMax )
//    {
//        iFreq++;
//    }
//
//    *fwhm = 2.0 * window->_freq[iFreq];
//    *iFwhm = iFreq;
//}
//
//
///**
// * @details
// * Find the index of the max value in an array
// *
// * @return The index of the max value
// */
//int Clean::_findMaxIndex(vector<Complex> values, unsigned iStart, unsigned iEnd)
//{
//    unsigned iMax = iStart;
//    double max = -1.0e99;
//
//    // loop through the array set the max index
//    for (unsigned int i = iStart; i <= iEnd; ++i)
//    {
//        double value = abs(values[i]);
//        if (value > max)
//        {
//            max = value;
//            iMax = i;
//        }
//    }
//
//    if (iMax == iStart || iMax == iEnd)
//        printf("WARNING: Clean::findMaxIndex() index at extreme of search range\n");
//
//    return iMax;
//}
//
///**
// * @details
// * Find the max in the array skipping a max at the bad index value
// *
// * @return the index of the max value
// */
//int Clean::_findMaxSafeIndex(std::vector<Complex> values, unsigned iStart,
//        unsigned iEnd, unsigned iBad)
//{
//    unsigned iMax = iStart;
//    double max = -1.0e99;
//
//    // loop through the array set the max index.
//    for (unsigned int i = iStart; i <= iEnd; ++i)
//    {
//        double value = abs(values[i]);
//
//        if (value > max && i != iBad)
//        {
//            max = value;
//            iMax = i;
//        }
//    }
//
//    if (iMax == iStart || iMax == iEnd)
//        printf("WARNING: Clean::findMaxIndex() index at extreme of search range\n");
//
//    return iMax;
//}
//
//
//

